"""
    EnKF Analysis Example

Comprehensive end-to-end EnKF workflow demonstration for the GSI/EnKF system.
This example shows realistic atmospheric data assimilation using ensemble Kalman filters
with comparison between pure EnKF, pure 3D-Var, and hybrid methods.

Features demonstrated:
- Complete EnKF workflow with multi-model support
- Realistic atmospheric data setup and ensemble generation  
- Performance benchmarking and timing analysis
- Diagnostic visualization and validation
- Comparison of different analysis methods
"""

using GSICoreAnalysis
using LinearAlgebra
using Random
using Statistics
using Printf
using Dates
using BenchmarkTools

"""
    setup_realistic_atmosphere()

Set up realistic atmospheric state and ensemble for testing.
"""
function setup_realistic_atmosphere()
    println("Setting up realistic atmospheric test case...")
    
    # Grid configuration (T574 resolution equivalent)
    nlons, nlats, nlevs = 1152, 576, 64
    ensemble_size = 40
    
    # Create realistic atmospheric fields
    # Longitude and latitude coordinates
    lons = range(-180.0, 180.0, length=nlons)
    lats = range(-90.0, 90.0, length=nlats)
    
    # Pressure levels (hPa)
    pressure_levels = [
        1000.0, 975.0, 950.0, 925.0, 900.0, 850.0, 800.0, 750.0, 700.0, 650.0,
        600.0, 550.0, 500.0, 450.0, 400.0, 350.0, 300.0, 250.0, 225.0, 200.0,
        175.0, 150.0, 125.0, 100.0, 70.0, 50.0, 30.0, 20.0, 10.0, 7.0,
        5.0, 3.0, 2.0, 1.0, 0.7, 0.5, 0.3, 0.2, 0.1, 0.07,
        0.05, 0.03, 0.02, 0.01, 0.007, 0.005, 0.003, 0.002, 0.001, 0.0007,
        0.0005, 0.0003, 0.0002, 0.0001, 0.00007, 0.00005, 0.00003, 0.00002, 0.00001, 0.000007,
        0.000005, 0.000003, 0.000002, 0.000001
    ][1:nlevs]
    
    # Initialize ensemble data structure
    ensemble_data = Dict{String,Any}()
    
    # Generate ensemble members
    ensemble_members = zeros(Float64, nlons, nlats, nlevs, ensemble_size)
    
    Random.seed!(12345)  # For reproducibility
    
    for member in 1:ensemble_size
        println(@sprintf("  Generating ensemble member %d/%d", member, ensemble_size))
        
        # Temperature field (realistic structure)
        for k in 1:nlevs
            temp_level = zeros(nlons, nlats)
            for j in 1:nlats
                lat = lats[j]
                # Realistic temperature profile
                temp_base = 288.15 - 6.5e-3 * (1000.0 - pressure_levels[k]) * 287.0 / 9.81
                temp_base += -40.0 * cos(lat * π / 180.0)^2  # Meridional gradient
                
                for i in 1:nlons
                    lon = lons[i]
                    # Add synoptic-scale variability
                    temp_level[i, j] = temp_base + 
                        5.0 * sin(3.0 * lon * π / 180.0) * cos(2.0 * lat * π / 180.0) +
                        randn() * 0.5  # Ensemble perturbation
                end
            end
            ensemble_members[1:nlons, 1:nlats, k, member] = temp_level
        end
    end
    
    ensemble_data["temperature"] = ensemble_members
    
    # Generate other meteorological variables (simplified for example)
    ensemble_data["u_wind"] = 0.1 * ensemble_members + randn(size(ensemble_members)) * 2.0
    ensemble_data["v_wind"] = 0.05 * ensemble_members + randn(size(ensemble_members)) * 1.5
    ensemble_data["specific_humidity"] = abs.(0.001 * ensemble_members + randn(size(ensemble_members)) * 0.0001)
    ensemble_data["surface_pressure"] = 101325.0 .+ randn(nlons, nlats, ensemble_size) * 500.0
    
    # Grid information
    grid_info = Dict(
        "nlons" => nlons,
        "nlats" => nlats,
        "nlevs" => nlevs,
        "lons" => lons,
        "lats" => lats,
        "pressure_levels" => pressure_levels,
        "ensemble_size" => ensemble_size
    )
    
    println("  Atmospheric test case setup complete")
    return ensemble_data, grid_info
end

"""
    generate_synthetic_observations(ensemble_data::Dict, grid_info::Dict)

Generate realistic synthetic observations for testing.
"""
function generate_synthetic_observations(ensemble_data::Dict, grid_info::Dict)
    println("Generating synthetic observations...")
    
    observations = Dict{String,Any}()
    
    nlons, nlats, nlevs = grid_info["nlons"], grid_info["nlats"], grid_info["nlevs"]
    lons, lats = grid_info["lons"], grid_info["lats"]
    
    # Surface pressure observations (weather stations)
    n_sfc_obs = 5000
    sfc_obs = Dict{String,Any}()
    sfc_obs["values"] = Float64[]
    sfc_obs["lons"] = Float64[]
    sfc_obs["lats"] = Float64[]
    sfc_obs["errors"] = Float64[]
    sfc_obs["type"] = "surface_pressure"
    
    Random.seed!(54321)
    for i in 1:n_sfc_obs
        # Random location (bias toward land areas)
        obs_lon = rand() * 360.0 - 180.0
        obs_lat = asin(2.0 * rand() - 1.0) * 180.0 / π
        
        # Find nearest grid point
        i_lon = argmin(abs.(lons .- obs_lon))
        j_lat = argmin(abs.(lats .- obs_lat))
        
        # Truth value (ensemble mean + small error)
        truth_value = mean(ensemble_data["surface_pressure"][i_lon, j_lat, :])
        obs_value = truth_value + randn() * 100.0  # 100 Pa observation error
        
        push!(sfc_obs["values"], obs_value)
        push!(sfc_obs["lons"], obs_lon)
        push!(sfc_obs["lats"], obs_lat)
        push!(sfc_obs["errors"], 100.0)
    end
    
    observations["surface_pressure"] = sfc_obs
    
    # Radiosonde observations
    n_rad_obs = 500
    rad_obs = Dict{String,Any}()
    rad_obs["values"] = Float64[]
    rad_obs["lons"] = Float64[]
    rad_obs["lats"] = Float64[]
    rad_obs["levels"] = Int[]
    rad_obs["errors"] = Float64[]
    rad_obs["type"] = "temperature"
    
    for i in 1:n_rad_obs
        obs_lon = rand() * 360.0 - 180.0
        obs_lat = (rand() - 0.5) * 160.0  # Avoid poles
        obs_level = rand(1:min(20, nlevs))  # Upper-air levels only
        
        i_lon = argmin(abs.(lons .- obs_lon))
        j_lat = argmin(abs.(lats .- obs_lat))
        
        truth_value = mean(ensemble_data["temperature"][i_lon, j_lat, obs_level, :])
        obs_value = truth_value + randn() * 1.0  # 1K observation error
        
        push!(rad_obs["values"], obs_value)
        push!(rad_obs["lons"], obs_lon)
        push!(rad_obs["lats"], obs_lat)
        push!(rad_obs["levels"], obs_level)
        push!(rad_obs["errors"], 1.0)
    end
    
    observations["radiosonde_temp"] = rad_obs
    
    # Satellite observations (simplified)
    n_sat_obs = 10000
    sat_obs = Dict{String,Any}()
    sat_obs["values"] = Float64[]
    sat_obs["lons"] = Float64[]
    sat_obs["lats"] = Float64[]
    sat_obs["channels"] = Int[]
    sat_obs["errors"] = Float64[]
    sat_obs["type"] = "satellite_radiance"
    
    for i in 1:n_sat_obs
        obs_lon = rand() * 360.0 - 180.0
        obs_lat = (rand() - 0.5) * 140.0
        channel = rand(1:10)
        
        i_lon = argmin(abs.(lons .- obs_lon))
        j_lat = argmin(abs.(lats .- obs_lat))
        
        # Simplified radiance calculation
        temp_weighting = exp.(-grid_info["pressure_levels"][1:min(20,nlevs)] / 200.0)
        temp_profile = mean(ensemble_data["temperature"][i_lon, j_lat, 1:length(temp_weighting), :], dims=2)[:, 1]
        radiance = sum(temp_weighting .* temp_profile)
        
        obs_value = radiance + randn() * 2.0  # 2K brightness temperature error
        
        push!(sat_obs["values"], obs_value)
        push!(sat_obs["lons"], obs_lon)
        push!(sat_obs["lats"], obs_lat)
        push!(sat_obs["channels"], channel)
        push!(sat_obs["errors"], 2.0)
    end
    
    observations["satellite"] = sat_obs
    
    println(@sprintf("  Generated %d surface pressure observations", length(sfc_obs["values"])))
    println(@sprintf("  Generated %d radiosonde temperature observations", length(rad_obs["values"])))
    println(@sprintf("  Generated %d satellite radiance observations", length(sat_obs["values"])))
    
    return observations
end

"""
    run_enkf_workflow_example()

Run complete EnKF workflow example with comprehensive diagnostics.
"""
function run_enkf_workflow_example()
    println("="^80)
    println("COMPREHENSIVE EnKF WORKFLOW EXAMPLE")
    println("="^80)
    
    # Setup
    ensemble_data, grid_info = setup_realistic_atmosphere()
    observations = generate_synthetic_observations(ensemble_data, grid_info)
    
    # Configuration
    config = AnalysisConfig(
        grid_size = (grid_info["nlons"], grid_info["nlats"], grid_info["nlevs"]),
        ensemble_size = grid_info["ensemble_size"],
        analysis_method = "EnKF",
        model_type = "GFS",
        output_path = "./enkf_output",
        model_params = Dict(
            "grid_type" => "gaussian",
            "reduced_grid" => false,
            "hybrid_vertical" => true,
            "file_format" => "netcdf",
            "variable_names" => ["temperature", "u_wind", "v_wind", "specific_humidity", "surface_pressure"]
        ),
        params = Dict(
            "localization_radius" => 200.0,  # km
            "inflation_factor" => 1.02,
            "log_level" => 2,
            "observation_types" => ["surface_pressure", "radiosonde_temp", "satellite"],
            "variable_names" => ["temperature", "u_wind", "v_wind", "specific_humidity"],
            "regional_analysis" => true,
            "latitude_bands" => [-90.0, -60.0, -30.0, 0.0, 30.0, 60.0, 90.0]
        )
    )
    
    println("\nConfiguration:")
    println("  Grid size: $(config.grid_size)")
    println("  Ensemble size: $(config.ensemble_size)")
    println("  Model type: $(config.model_type)")
    println("  Analysis method: $(config.analysis_method)")
    
    # Create output directory
    mkpath(config.output_path)
    
    # Run EnKF Analysis
    println("\n" * "="^60)
    println("RUNNING EnKF ANALYSIS")
    println("="^60)
    
    start_time = now()
    enkf_result = run_enkf_analysis(ensemble_data, observations, config)
    enkf_time = (now() - start_time).value / 1000.0
    
    println(@sprintf("EnKF analysis completed in %.2f seconds", enkf_time))
    
    return enkf_result, config, (enkf_time, 0.0, 0.0)
end

"""
    run_method_comparison()

Compare EnKF, 3D-Var, and Hybrid methods on the same data.
"""
function run_method_comparison()
    println("\n" * "="^60)
    println("METHOD COMPARISON: EnKF vs 3D-Var vs Hybrid")
    println("="^60)
    
    # Setup common data
    ensemble_data, grid_info = setup_realistic_atmosphere()
    observations = generate_synthetic_observations(ensemble_data, grid_info)
    
    results = Dict{String,Any}()
    timings = Dict{String,Float64}()
    
    # EnKF Analysis
    println("\n--- Running EnKF Analysis ---")
    config_enkf = AnalysisConfig(
        grid_size = (grid_info["nlons"], grid_info["nlats"], grid_info["nlevs"]),
        ensemble_size = grid_info["ensemble_size"],
        analysis_method = "EnKF",
        model_type = "GFS",
        output_path = "./comparison_output/enkf",
        params = Dict("log_level" => 1)
    )
    mkpath(config_enkf.output_path)
    
    start_time = now()
    results["EnKF"] = run_enkf_analysis(ensemble_data, observations, config_enkf)
    timings["EnKF"] = (now() - start_time).value / 1000.0
    
    # 3D-Var Analysis
    println("\n--- Running 3D-Var Analysis ---")
    config_3dvar = AnalysisConfig(
        grid_size = (grid_info["nlons"], grid_info["nlats"], grid_info["nlevs"]),
        analysis_method = "3DVar",
        output_path = "./comparison_output/3dvar",
        params = Dict("log_level" => 1)
    )
    mkpath(config_3dvar.output_path)
    
    start_time = now()
    results["3DVar"] = run_analysis("3DVar", ensemble_data, observations, config_3dvar)
    timings["3DVar"] = (now() - start_time).value / 1000.0
    
    # Hybrid Analysis
    println("\n--- Running Hybrid Analysis ---")
    config_hybrid = AnalysisConfig(
        grid_size = (grid_info["nlons"], grid_info["nlats"], grid_info["nlevs"]),
        ensemble_size = grid_info["ensemble_size"],
        analysis_method = "Hybrid",
        model_type = "GFS",
        output_path = "./comparison_output/hybrid",
        params = Dict(
            "log_level" => 1,
            "static_covar_weight" => 0.3,
            "ensemble_covar_weight" => 0.7,
            "variational_method" => "3DVar"
        )
    )
    mkpath(config_hybrid.output_path)
    
    start_time = now()
    control_vector = Dict("analysis" => mean(ensemble_data["temperature"], dims=4)[:,:,:,1])
    results["Hybrid"] = run_hybrid_analysis(control_vector, ensemble_data, observations, config_hybrid)
    timings["Hybrid"] = (now() - start_time).value / 1000.0
    
    # Print comparison summary
    println("\n" * "="^60)
    println("METHOD COMPARISON SUMMARY")
    println("="^60)
    
    println("Runtime Comparison:")
    for method in ["EnKF", "3DVar", "Hybrid"]
        println(@sprintf("  %-10s: %8.2f seconds", method, timings[method]))
    end
    
    println("\nAnalysis Quality Metrics:")
    for method in ["EnKF", "3DVar", "Hybrid"]
        if haskey(results[method], "cost_function_value")
            println(@sprintf("  %-10s Cost Function: %12.6f", method, results[method]["cost_function_value"]))
        end
    end
    
    # Save comparison results
    comparison_file = "./comparison_output/method_comparison_summary.txt"
    mkpath(dirname(comparison_file))
    open(comparison_file, "w") do f
        println(f, "GSI/EnKF Method Comparison Summary")
        println(f, "=" ^ 50)
        println(f, "Generated: ", Dates.format(now(), "yyyy-mm-dd HH:MM:SS"))
        println(f, "")
        
        println(f, "Runtime Comparison:")
        for method in ["EnKF", "3DVar", "Hybrid"]
            println(f, @sprintf("  %-10s: %8.2f seconds", method, timings[method]))
        end
        
        println(f, "")
        println(f, "Observation Processing:")
        println(f, @sprintf("  Surface pressure obs: %d", length(observations["surface_pressure"]["values"])))
        println(f, @sprintf("  Radiosonde obs:       %d", length(observations["radiosonde_temp"]["values"])))
        println(f, @sprintf("  Satellite obs:        %d", length(observations["satellite"]["values"])))
    end
    
    println("Comparison results saved to: $comparison_file")
    return results, timings
end

"""
    benchmark_performance()

Perform detailed performance benchmarking of EnKF components.
"""
function benchmark_performance()
    println("\n" * "="^60)
    println("PERFORMANCE BENCHMARKING")
    println("="^60)
    
    # Setup smaller test case for benchmarking
    println("Setting up benchmark test case...")
    nlons, nlats, nlevs = 288, 144, 32  # Smaller grid
    ensemble_size = 20
    
    # Generate test data
    ensemble_test = Dict{String,Any}()
    ensemble_test["temperature"] = randn(nlons, nlats, nlevs, ensemble_size)
    ensemble_test["u_wind"] = randn(nlons, nlats, nlevs, ensemble_size)
    ensemble_test["surface_pressure"] = randn(nlons, nlats, ensemble_size)
    
    # Simple observations
    observations_test = Dict{String,Any}()
    observations_test["surface_pressure"] = Dict(
        "values" => randn(1000),
        "lons" => rand(1000) * 360.0 .- 180.0,
        "lats" => rand(1000) * 160.0 .- 80.0,
        "errors" => fill(100.0, 1000)
    )
    
    config_bench = AnalysisConfig(
        grid_size = (nlons, nlats, nlevs),
        ensemble_size = ensemble_size,
        analysis_method = "EnKF",
        model_type = "GFS",
        output_path = "./benchmark_output",
        params = Dict("log_level" => 0)
    )
    mkpath(config_bench.output_path)
    
    # Benchmark different components
    println("Benchmarking EnKF components...")
    
    # Model interface initialization
    println("  Model interface initialization...")
    model_init_time = @benchmark initialize_model_interface("GFS", $config_bench.model_params) seconds=10
    
    # Ensemble spread calculation
    println("  Ensemble spread calculation...")
    spread_time = @benchmark calculate_ensemble_spread($ensemble_test["temperature"], ["temperature"]) seconds=10
    
    # Memory usage monitoring
    println("  Memory usage monitoring...")
    diag_config = DiagnosticConfig(
        output_path = config_bench.output_path,
        log_level = 0,
        save_statistics = false,
        generate_plots = false,
        observation_types = ["surface_pressure"],
        variable_names = ["temperature"],
        regional_analysis = false,
        latitude_bands = [-90.0, 0.0, 90.0],
        quality_control_thresholds = Dict{String,Float64}(),
        performance_monitoring = true,
        memory_tracking = true
    )
    
    memory_time = @benchmark track_memory_usage($diag_config) seconds=10
    
    # Print benchmark results
    println("\nBenchmark Results:")
    println(@sprintf("  Model interface init:   %8.3f ± %6.3f ms", 
            mean(model_init_time.times) / 1e6, std(model_init_time.times) / 1e6))
    println(@sprintf("  Ensemble spread calc:   %8.3f ± %6.3f ms", 
            mean(spread_time.times) / 1e6, std(spread_time.times) / 1e6))
    println(@sprintf("  Memory usage monitor:   %8.3f ± %6.3f ms", 
            mean(memory_time.times) / 1e6, std(memory_time.times) / 1e6))
    
    # Save benchmark results
    benchmark_file = "./benchmark_output/performance_benchmark.txt"
    open(benchmark_file, "w") do f
        println(f, "GSI/EnKF Performance Benchmark Results")
        println(f, "=" ^ 50)
        println(f, "Generated: ", Dates.format(now(), "yyyy-mm-dd HH:MM:SS"))
        println(f, "")
        println(f, "Test Configuration:")
        println(f, @sprintf("  Grid size: %d × %d × %d", nlons, nlats, nlevs))
        println(f, @sprintf("  Ensemble size: %d", ensemble_size))
        println(f, @sprintf("  Observations: %d", length(observations_test["surface_pressure"]["values"])))
        println(f, "")
        println(f, "Component Timings:")
        println(f, @sprintf("  Model interface init:   %8.3f ± %6.3f ms", 
                mean(model_init_time.times) / 1e6, std(model_init_time.times) / 1e6))
        println(f, @sprintf("  Ensemble spread calc:   %8.3f ± %6.3f ms", 
                mean(spread_time.times) / 1e6, std(spread_time.times) / 1e6))
        println(f, @sprintf("  Memory usage monitor:   %8.3f ± %6.3f ms", 
                mean(memory_time.times) / 1e6, std(memory_time.times) / 1e6))
    end
    
    println("Benchmark results saved to: $benchmark_file")
    
    return Dict(
        "model_init" => model_init_time,
        "spread_calc" => spread_time,
        "memory_monitor" => memory_time
    )
end

"""
    main()

Run all examples and demonstrations.
"""
function main()
    println("GSI/EnKF Julia Package - Comprehensive Example Suite")
    println("=" ^ 80)
    
    try
        # Run main workflow example
        println("\n1. Running comprehensive EnKF workflow example...")
        enkf_result, config, timings = run_enkf_workflow_example()
        
        # Run method comparison
        println("\n2. Running method comparison...")
        comparison_results, comparison_timings = run_method_comparison()
        
        # Run performance benchmark
        println("\n3. Running performance benchmark...")
        benchmark_results = benchmark_performance()
        
        # Final summary
        println("\n" * "="^80)
        println("EXAMPLE SUITE COMPLETION SUMMARY")
        println("="^80)
        
        println("All examples completed successfully!")
        println("\nOutput directories created:")
        println("  ./enkf_output/              - EnKF workflow results")
        println("  ./comparison_output/        - Method comparison results")
        println("  ./benchmark_output/         - Performance benchmark results")
        
        println("\nKey files generated:")
        println("  enkf_diagnostic_report.txt  - Comprehensive EnKF diagnostics")
        println("  method_comparison_summary.txt - Method comparison summary")
        println("  performance_benchmark.txt    - Performance benchmark results")
        
        println("\nNext steps:")
        println("  - Review diagnostic reports for analysis quality assessment")
        println("  - Examine performance benchmarks for optimization opportunities")
        println("  - Customize configurations for your specific applications")
        println("  - Integrate with operational forecast systems")
        
    catch e
        println("ERROR: Example execution failed with exception:")
        println(e)
        rethrow(e)
    end
    
    return true
end

# Run examples if script is executed directly
if abspath(PROGRAM_FILE) == @__FILE__
    main()
end